
import os
from ray.data import read_datasource, ActorPoolStrategy

from s3_datasource import S3Datasource
from process_pdf_enroll import BatchDocInfer



source_prefix = "s3://llm-process-pperf/ebook_index_v4/ebook/v006/zhongwenzaixian/pdf/"
target_prefix = "s3://llm-pdf-text-1/pdf_gpu_output/zhongwenzaixian/v006/"


class ProcArgs:
    def __init__(self):
        self.infer = BatchDocInfer(True, True)

    def __call__(self, batch):
        for fn in batch["s3path"]:
            self.infer(fn, target_prefix + os.path.basename(fn))
        return {"results": []}


if __name__ == "__main__":
    import ray
    runtime_env = {"working_dir": "/cpfs01/user/xurui/doc-infer/ray-pipeline",}

    ray.init(runtime_env=runtime_env)
    ds = S3Datasource(source_prefix, target_prefix)
    dataset = read_datasource(ds, parallelism=20).map_batches(ProcArgs, num_gpus=1, batch_size=1, compute=ActorPoolStrategy(size=20)).take_all()
